package com.dmai.aiet.common.filter;

import com.dmai.aiet.common.properties.TraceProperties;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.Charset;

@Slf4j
@Configuration
@ConditionalOnProperty(value = "app.trace.enabled", havingValue = "true", matchIfMissing = true)
public class ResponseLogFilter implements WebFilter, Ordered {

    @Autowired
    private TraceProperties traceProperties;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {

        //获取response的 返回数据
        ServerHttpResponse originalResponse = exchange.getResponse();
        DataBufferFactory bufferFactory = originalResponse.bufferFactory();
        ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(originalResponse) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
//                MDC.put(MDCTraceUtils.KEY_TRACE_ID, exchange.getRequest().getHeaders().getFirst(MDCTraceUtils.TRACE_ID_HEADER));
//                MDC.put(MDCTraceUtils.KEY_SPAN_ID, exchange.getRequest().getHeaders().getFirst(MDCTraceUtils.SPAN_ID_HEADER));
                if (traceProperties.getResponseData() && getStatusCode().equals(HttpStatus.OK) && body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                    return super.writeWith(fluxBody.map(dataBuffer -> {
                        byte[] content = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(content);
                        //释放掉内存
                        DataBufferUtils.release(dataBuffer);

                        String responseData = new String(content, Charset.forName("UTF-8"));

                        log.info("response：{}", responseData);

                        return bufferFactory.wrap(content);
                    }));
                } else {
                    log.info("response：{}", getStatusCode());
                }
                return super.writeWith(body);
            }
        };
        return chain.filter(exchange.mutate().response(decoratedResponse).build()).then(Mono.fromRunnable(() -> {
            // 如果请求被限流，记录日志
//            MDC.put(MDCTraceUtils.KEY_TRACE_ID, exchange.getRequest().getHeaders().getFirst(MDCTraceUtils.TRACE_ID_HEADER));
//            MDC.put(MDCTraceUtils.KEY_SPAN_ID, exchange.getRequest().getHeaders().getFirst(MDCTraceUtils.SPAN_ID_HEADER));
            if (exchange.getResponse().getStatusCode() == HttpStatus.TOO_MANY_REQUESTS) {
                log.warn("response：{} flow limited from client {}",
                        exchange.getResponse().getStatusCode(), exchange.getRequest().getRemoteAddress());
            }
        }));

    }

    @Override
    public int getOrder() {
        return -2;
    }
}
